/*
 * Copyright (c) 2024 iSOFT INFRASTRUCTURE SOFTWARE CO., LTD.
 * easyAda is licensed under Mulan PubL v2.
 * You can use this software according to the terms and conditions of the Mulan PubL v2.
 * You may obtain a copy of Mulan PubL v2 at:
 *          http://license.coscl.org.cn/MulanPubL-2.0
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PubL v2 for more details.
 */

#include <plat/config.h>
#include <arch/utils.h>
#include <stdlib/assert.h>
#include <stdlib/stdint.h>
#include <tools/macros.h>
#include <tools/rbtree.h>
#include <core/pbuddy.h>
#include <core/addrspace.h>
#include <core/allocate.h>
#include <core/dprintf.h>
#include <core/object.h>
#include <caps/ram.h>
#include <uapi/bootinfo.h>

extern const region_t   ksram[];
extern const region_t   paksram[];
extern const region_t   usram[];
unsigned int fragment;

spinlock_t pgtbl_bitmap_lock = SPINLOCK_INIT_UNLOCKED;

#define BITMAP_FULL_PAGE       0xFFFFFFFFUL
#define PAGE_TABLE_BIT_SHIFT   5
#define PAGE_TABLE_BITMAP_SIZE ((CONFIG_PAGE_TABLE_RAM_SIZE >> CONFIG_PAGE_SHIFT) >> PAGE_TABLE_BIT_SHIFT)
unsigned int pgtbl_bitmap[PAGE_TABLE_BITMAP_SIZE];

volatile static unsigned int allocator_all_pages;
volatile static unsigned int allocator_free_pages;

static int paksram_initd = 0;
#define INIT_APACES 30
static unsigned char initksram_apace[INIT_APACES][CONFIG_PAGE_SIZE] ALIGNED(CONFIG_PAGE_SIZE);

static struct stack_ipcbuffer *__pinit_si;

pbuddy_t *p__pbuddy = NULL;

static inline unsigned int size_align(unsigned int size)
{
    fragment -= size;
    size = (size + 0x7) & ~(0x7U);
    fragment += size;
    return size;
}

extern int boot_init_flag;
unsigned long allocator_pbase(unsigned long vbase)
{
    if (boot_init_flag == 1) {
        return vbase;
    }  
#ifdef DISABLE_BUDDY_PAGE_TABLE
    if (vaddr_in_region(vbase, paksram) != 0) {
        return paksram->pbase + (vbase - paksram->vbase);
    }

    if (vaddr_in_region(vbase, ksram) != 0) {
        return ksram->pbase + (vbase - ksram->vbase);
    }

    return 0;
#else
    unsigned long v_begin;
    unsigned long v_end;
    int           i = 0;

    while (ksram[i].size > 0) {
        v_begin = ksram[i].vbase;
        v_end   = v_begin + ksram[i].size;
        if (vbase >= v_begin && vbase < v_end) {
            return vbase - v_begin + ksram[i].pbase;
        }
        i++;
    }

    i = 0;
    while (paksram[i].size > 0) {
        v_begin = paksram[i].vbase;
        v_end   = v_begin + paksram[i].size;
        if (vbase >= v_begin && vbase < v_end) {
            return vbase - v_begin + paksram[i].pbase;
        }
        i++;
    }

    return 0;
#endif
}

unsigned long allocator_vbase(unsigned long pbase)
{
    if (boot_init_flag == 1) {
        return pbase;
    }  

#ifdef DISABLE_BUDDY_PAGE_TABLE

    if (paddr_in_region(pbase, ksram) != 0) {
        return ksram->vbase + (pbase - ksram->pbase);
    }
    if (paddr_in_region(pbase, paksram) != 0) {
        return paksram->vbase + (pbase - paksram->pbase);
    }
    return 0;

#else
    unsigned long p_begin;
    unsigned long p_end;
    int           i = 0;

    while (ksram[i].size > 0) {
        p_begin = ksram[i].pbase;
        p_end   = p_begin + ksram[i].size;
        if (pbase >= p_begin && pbase < p_end) {
            return pbase - p_begin + ksram[i].vbase;
        }
        i++;
    }

    i = 0;
    while (paksram[i].size > 0) {
        p_begin = paksram[i].pbase;
        p_end   = p_begin + paksram[i].size;
        if (pbase >= p_begin && pbase < p_end) {
            return pbase - p_begin + paksram[i].vbase;
        }
        i++;
    }

    return 0;
#endif
}

static void shrink_region(region_t *curr, const region_t *reserved, int count)
{
    int rend;
    int bend;
    int j;

    for (j = 0; j < count; j++, reserved++) {
        rend = reserved->pbase + reserved->size;
        bend = curr->pbase + curr->size;

        if (reserved->pbase <= curr->pbase) {
            if (rend > curr->pbase) {
                if (rend >= bend) {
                    curr->pbase = bend;
                    curr->size  = 0;
                    break;
                } else {
                    curr->pbase = rend;
                    curr->size  = bend - rend;
                }
            }
        } else if (reserved->pbase < bend) {
            curr->size = reserved->pbase - curr->pbase;
        } else {
            continue;
        }
    }
}

static void add_region(region_t *remain, uint64_t *initsect_v, uint64_t *initsect_f)
{
    region_t curr;
    int      seci = 0;

    assert((remain->pbase & 0x7) == 0);
    assert((remain->vbase & 0x7) == 0);

    assert((remain->pbase & lowbitsmask(CONFIG_PAGE_SHIFT)) == 0);
    assert((remain->vbase & lowbitsmask(CONFIG_PAGE_SHIFT)) == 0);

    while (remain->size >= sizeof(void *)) {
        void *vbase_0 = (void *)remain->vbase;
        int   lbsize  = biol1(remain->size);
        assert(lbsize >= CONFIG_PAGE_SHIFT);
        assert(1 << lbsize == remain->size);
        curr.pbase = remain->pbase;
        curr.size  = remain->size;
        curr.vbase = remain->vbase;
        shrink_region(&curr, (region_t *)&kernelinfo, sizeof(image_info_t) / sizeof(region_t));
        shrink_region(&curr, (region_t *)&rootserverinfo, sizeof(image_info_t) / sizeof(region_t));

        if ((curr.pbase & 0x7) > 0) {
            curr.size -= 8 - (curr.pbase & 0x7);
            fragment += 8 - (curr.pbase & 0x7);
        }
        curr.size  = size_align(curr.size);
        curr.pbase = align_to(curr.pbase, 8);
        curr.vbase += curr.pbase - remain->pbase;
        remain->size  = remain->pbase + remain->size - (curr.pbase + curr.size);
        remain->pbase = curr.pbase + curr.size;
        remain->vbase = curr.vbase + curr.size;

        if (curr.size <= sizeof(void *)) {
            fragment += curr.size;
            continue;
        }
        if (curr.size % CONFIG_PAGE_SIZE) {
            fragment += (curr.size - curr.size / CONFIG_PAGE_SIZE * CONFIG_PAGE_SIZE);
            curr.size = curr.size / CONFIG_PAGE_SIZE * CONFIG_PAGE_SIZE;
        }
        assert(lbsize < 32);
        assert(seci < MAX_RAM_INFO);
        initsect_f[seci] |= FLAG_RAM_KERNEL;
        initsect_f[seci] <<= 8;
        initsect_f[seci] |= lbsize;
        initsect_f[seci] <<= 32;
        initsect_f[seci] |= curr.size;
        initsect_v[seci] = (uint64_t)vbase_0;
        seci++;
    }
    fragment += remain->size;
    return;
}

void *get_rs_stack(void)
{
    return (void *)&__pinit_si->stack;
}

bootinfo_t *get_rs_bootinfo(void)
{
    return (bootinfo_t *)&__pinit_si->bi;
}

void *get_rs_ipcbuffer(void)
{
    return (void *)&__pinit_si->buf;
}

void *get_rs_faulthandler(void)
{
    return (void *)&__pinit_si->faulthandler;
}

void allocator_init()
{
    __pinit_si = (void *) CONFIG_RS_VBASE + CONFIG_RS_SIZE - sizeof(struct stack_ipcbuffer);

    allocator_all_pages  = CONFIG_PAGE_TABLE_RAM_SIZE >> CONFIG_PAGE_SHIFT;
    allocator_free_pages = CONFIG_PAGE_TABLE_RAM_SIZE >> CONFIG_PAGE_SHIFT;
}

int check_usram_region(void *begin, unsigned int size)
{
    return 1;
}

void *allocator_alloc_block(unsigned int size)
{
    void *result = (void *)0;
#ifdef DISABLE_BUDDY_PAGE_TABLE

    assert(size == PAGE_SIZE);
    if (paksram_initd < INIT_APACES) {
        paksram_initd++;
        return (void *)initksram_apace[paksram_initd - 1];
    }

    spin_lock(&pgtbl_bitmap_lock);
    for (int i = 0; i < PAGE_TABLE_BITMAP_SIZE; i++) {
        if (pgtbl_bitmap[i] != BITMAP_FULL_PAGE) {
            int pos = get_bit0_pos(pgtbl_bitmap[i]);
            set_pos_bit1(&pgtbl_bitmap[i], pos);
            unsigned int pfn = ((i << PAGE_TABLE_BIT_SHIFT) + pos);
            allocator_free_pages--;
            result = (void *)(paksram->vbase + (pfn << CONFIG_PAGE_SHIFT));
            break;
        }
    }
    spin_unlock(&pgtbl_bitmap_lock);
    return result;
#else
    if (!size || (size > 0x80000000U)) {
        return 0;
    }
    extern const region_t paksram[];
    if (paksram_initd < INIT_APACES) {
        assert(CONFIG_PAGE_SIZE == size);
        paksram_initd++;
        return (void *)initksram_apace[paksram_initd - 1];
    } else {
        int bsize = (1 << biol1(size));
        bsize     = (bsize == size) ? size : (bsize * 2);
        ppage_t *rtn;
        if (pbuddy_alloc_page(p__pbuddy, bsize >> CONFIG_PAGE_SHIFT, &rtn)) {
            return 0;
        }
        return rtn->base;
    }
    return 0;
#endif
}

int allocator_free_block(unsigned long begin)
{
#ifdef DISABLE_BUDDY_PAGE_TABLE

    if (vaddr_in_region(begin, paksram) != 0) {
        unsigned int pfn        = (begin - paksram->vbase) >> CONFIG_PAGE_SHIFT;
        unsigned int bitmap_pos = pfn >> PAGE_TABLE_BIT_SHIFT;
        unsigned int bit_pos    = pfn % (1 << PAGE_TABLE_BIT_SHIFT);
        spin_lock(&pgtbl_bitmap_lock);
        unsigned int *value_p   = pgtbl_bitmap + bitmap_pos;
        unsigned int  value_src = *(value_p);
        set_pos_bit0(pgtbl_bitmap + bitmap_pos, bit_pos);
        if (value_src ^ *value_p) {
            allocator_free_pages++;
        } else {
            dprintf("error free ==> the addr is free \n");
        }
        spin_unlock(&pgtbl_bitmap_lock);
        return 0;
    }
    if (vaddr_in_region(begin, ksram) != 0) {
        return 0;
    }
    return -1;
#else
    ppage_t *p;
    if (pbuddy_find_page(p__pbuddy, (void *)begin, &p)) {
        return -1;
    }
    assert(p->flags);
    if ((void *)begin == p->base) {
        pbuddy_free_page(p__pbuddy, p);
    }
    return 0;
#endif
}

int allocator_free(unsigned long begin, unsigned int size)
{
    dprintf(" *** mem leak! [allocator_free]***\n");
    return -1;
}

void show_pagetbale_alloctor_status(void)
{
    unsigned int a_f_p = allocator_free_pages;
    unsigned int a_a_p = allocator_all_pages;
    dprintf("free/all => %d / %d \n", a_f_p, a_a_p);
    for (int i = 0; i < PAGE_TABLE_BITMAP_SIZE; i++) {
        if (pgtbl_bitmap[i] != (unsigned int)0) {
            dprintf("bitmap[%d].addr = %x value  =  %x \n", i, &pgtbl_bitmap[i], pgtbl_bitmap[i]);
        }
    }
}

void pagetable_alloctor_test(void)
{
    unsigned long alloc_size = 0;
    void         *va0;

    while (alloc_size < (128 * 1024 * 1024 - 4096 * 64)) {
        va0 = allocator_alloc_block(PAGE_SIZE);

        if (va0 != NULL) {
            *(unsigned long *)va0 = 0x123456;
            alloc_size += PAGE_SIZE;
        }
    }

    dprintf("alloc all page bitmap status ==> \n");
    for (size_t i = 0; i < PAGE_TABLE_BITMAP_SIZE; i++) {
        dprintf("bitmap[%d] =  %x \n", i, pgtbl_bitmap[i]);
    }

    while (alloc_size != 0) {
        (void)allocator_free_block((unsigned long)va0);
        va0 -= PAGE_SIZE;
        alloc_size -= PAGE_SIZE;
    }

    dprintf("finish the pagetable alloc test. check pagetable_bitmap status:\n");
    for (size_t i = 0; i < PAGE_TABLE_BITMAP_SIZE; i++) {
        dprintf("bitmap[%d] =  %x \n", i, pgtbl_bitmap[i]);
    }

    dprintf("finsih the page_alloc test. \n");
}
